import re
import os

import pandas as pd
from typing import Literal, Dict, Any
from py2neo import Graph
from langchain_openai import ChatOpenAI

from typing import TypedDict, Annotated
import operator
from IPython.display import Image

from langchain_core.messages import (
    AnyMessage,
    SystemMessage,
    HumanMessage,
    ToolMessage,
    AIMessage,
)
from langgraph.graph import StateGraph, END, START

CYPHER_WRITER_PROMPT = """
你是一位知识图谱专家，擅长编写 Cypher 查询语句。你的任务是根据用户的需求，利用下述 `Company` 节点的属性和别名，生成精准的 Cypher 查询语句。

以下是 `Company` 节点的属性及其对应的别名（用户可能使用这些别名来描述需求）：

- `CompanyName` -> 企业名称
- `EnglishName` -> 英文名称
- `UnifiedSocialCreditCode` -> 统一社会信用代码
- `CompanyType` -> 企业类型
- `OperatingStatus` -> 经营状态
- `EstablishmentDate` -> 成立日期
- `ApprovalDate` -> 核准日期
- `LegalRepresentative` -> 法定代表人
- `RegisteredCapital` -> 注册资本
- `PaidUpCapital` -> 实缴资本
- `InsuredEmployees` -> 参保人数
- `CompanySize` -> 公司规模
- `BusinessScope` -> 经营范围
- `RegisteredAddress` -> 注册地址
- `BusinessTerm` -> 营业期限
- `TaxpayerIdentificationNumber` -> 纳税人识别号
- `BusinessRegistrationNumber` -> 工商注册号
- `OrganizationCode` -> 组织机构代码
- `TaxpayerQualification` -> 纳税人资质
- `FormerName` -> 曾用名
- `Province` -> 所属省份
- `City` -> 所属城市
- `District` -> 所属区县
- `WebsiteLink` -> 网站链接
- `Industry` -> 所属行业
- `PrimaryIndustryCategory` -> 一级行业分类
- `RegistrationAuthority` -> 登记机关
- `Longitude` -> 经度
- `Latitude` -> 纬度
- `WebsiteURL` -> 网址
- `HydrogenIndustryChainNode` -> 氢能产业链节点
- `ReasonForInclusionInChain` -> 上链原因

其中HydrogenIndustryChainNode包含：制氢、储氢和运输、加氢站、交通领域、工业领域、发电领域、建筑领域

请使用这些属性和对应的别名，依据用户的具体需求，生成最为合适的 Cypher 查询语句， 生成的每个cypher查询语句都以;结尾。
""".lstrip()

from langgraph.checkpoint.memory import MemorySaver

memory = MemorySaver()


class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]  # 消息列表
    query: str  # 当前查询
    cyphers: list[str]  # 当前全部的cypher查询语句
    cur_query_idx: int  # 当前cypher的索引
    cur_cypher: str  # 当前cypher查询语句, 因为有modify_cypher, 所以需要保存
    except_cnt: int  # 异常次数, 超过3次则停止
    cypher_result: list
    modify_cypher: str  # 修改后的cypher
    cypher_err: str  # cypher执行错误


class Neo4jAgent:
    def __init__(self, checkpointer=None, examples=[]):
        self.examples = examples
        graph = StateGraph(AgentState)

        graph.set_entry_point("human_input")
        graph.add_node("human_input", self.human_input)
        graph.add_node("neo4j_chat", self.neo4j_chat)
        graph.add_node("other_chat", self.other_chat)
        graph.add_node("run_cypher", self.run_cypher)
        graph.add_node("modify_cypher", self.modify_cypher)
        graph.add_node("cypher_output", self.cypher_output)

        graph.add_edge("neo4j_chat", "run_cypher")

        # graph.add_edge("run_cypher", "human_input")

        graph.add_edge("other_chat", "human_input")
        graph.add_edge("modify_cypher", "run_cypher")

        graph.add_conditional_edges(
            "human_input",
            self.router,
            {"neo4j_chat": "neo4j_chat", "other_chat": "other_chat", "__end__": END},
        )

        graph.add_conditional_edges(
            "run_cypher",
            self.run_cypher_router,
            {"cypher_output": "cypher_output", "modify_cypher": "modify_cypher", "human_input": "human_input"},
        )

        graph.add_conditional_edges(
            "cypher_output",
            self.cypher_output_router,
            {"run_cypher": "run_cypher", "human_input": "human_input"},
        )

        self.graph = graph.compile(checkpointer=checkpointer)

        self.llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
        self.driver = Graph("bolt://localhost:7687", auth=("neo4j", "12345678"))

    def human_input(self, state: AgentState):
        message = input("Human> ")
        # TODO: 默认存储的就是氢能企业，为了大模型在属性里匹配氢能，故目前暂定删去用户输入中的氢能。
        message = message.replace("氢能", "")
        return {
            "messages": [HumanMessage(message)],
            "query": message,
            "cyphers": [],
            "except_cnt": 0,
            "cur_query_idx": 0,
            "cypher_result": [],
            "modify_cypher": "",
            "cur_cypher": "",
        }

    def neo4j_chat(self, state: AgentState):
        """
        根据用户输入，生成cypher查询语句
        """
        content = state["messages"][-1].content
        # TODO： cypher_query_generator, 能够提取出多个cypher，并返回一个列表

        global CYPHER_WRITER_PROMPT
        response = self.llm.invoke(CYPHER_WRITER_PROMPT + content).content
        cyphers = re.findall(r"match(.+?);", response, re.DOTALL)
        return {
            "cyphers": cyphers,
            "except_cnt": 0,
            "cur_query_idx": 0,
        }

    def run_cypher(self, state: AgentState):
        """
        遇到异常让大模型重新生成，如果一直异常怎么办？一直浪费token? 必须限制
        """
        cyphers = state["cyphers"]
        except_cnt = state["except_cnt"]
        cur_query_idx = state["cur_query_idx"]
        cur_cypher = cyphers[cur_query_idx]
        query = state["query"]

        modify_cypher = state["modify_cypher"]
        if modify_cypher:
            cur_cypher = modify_cypher

        try:
            res = self.driver.run(cur_cypher)
            return {
                "cypher_result": res,
                "modify_cypher": "",  # 清空 modify_cypher
                "cur_cypher": cur_cypher,
                "cur_query_idx": cur_query_idx + 1,
                "except_cnt": 0,
            }

        except Exception as e:
            # 提交错误信息给 modify_cypher
            if except_cnt < 3:
                return {
                    "modify_cypher": cur_cypher,
                    "cypher_err": str(e),
                    "except_cnt": except_cnt + 1
                }
            else:
                # 如果连续3次异常，则返回无法解决该问题, back to human_input
                response = f"在解决该问题: {query}。\n出现了以下错误：{str(e)}，\n在经过3次尝试后，依然无法解决该问题。"
                return {
                    "messages": [AIMessage(content=response)],
                    # "except_cnt": 0, # 重置异常计数
                    "except_cnt": except_cnt + 1,
                    "modify_cypher": "",  # 清空 modify_cypher, give up
                }

    def modify_cypher(self, state: AgentState) -> dict[str, str]:
        modify_cypher = state["modify_cypher"]
        cypher_err = state["cypher_err"]

        MODIFY_PROMPT = "在运行cypher语句时出现了异常，请修改cypher语句，使其能够正常运行。异常信息为：{cypher_err}，cypher语句为：{modify_cypher}。"
        content = self.llm.invoke(MODIFY_PROMPT.format(cypher_err=cypher_err, modify_cypher=modify_cypher)).content
        cypher = re.findall(r"match(.+?);", content, re.DOTALL)[0]
        return {
            "modify_cypher": cypher
        }

    def router(self, state) -> Literal["neo4j_chat", "other_chat", "__end__"]:

        content = state["messages"][-1].content
        keywords = ["产业", "企业", "公司", "数量"]

        for kw in keywords:
            if kw in content:
                return "neo4j_chat"

        if "q" == content or "exit" == content:
            return "__end__"

        return "other_chat"

    def run_cypher_router(self, state: AgentState) -> Literal["cypher_output", "modify_cypher", "human_input"]:
        cypher_result = state.get("cypher_result", [])
        modify_cypher = state.get("modify_cypher", "")
        if cypher_result:
            return "cypher_output"
        elif modify_cypher:
            "modify_cypher"
        else:
            return "human_input"

    def cypher_output_router(self, state: AgentState) -> Literal["run_cypher", "human_input"]:
        cur_query_idx = state["cur_query_idx"]
        if cur_query_idx < len(state["cyphers"]):
            return "run_cypher"
        return "human_input"

    def other_chat(self, state: AgentState):
        """
        换成用户有输入，当用户输入 q、exit 退出
        * 搜索引擎
        * 图数据库
        * 退出
        条件怎么写？
        """
        # messages = [SystemMessage(self.system)] + self.examples + state["messages"]
        content = state["messages"][-1].content
        print(content)
        response = self.llm.invoke(content).content
        return {"messages": [AIMessage(content=response)]}

    def cypher_output(self, state: AgentState):
        """
        * 如果只有一条结果，直接返回结果；
        * 如果有多个结果，根据用户的问题和cypher，创建文件名，保存结果；
        """
        query = state["query"]
        cypher_results = state["cypher_result"]

        cyphers = state["cyphers"]
        except_cnt = state["except_cnt"]
        cur_query_idx = state["cur_query_idx"]
        cypher_query = cyphers[cur_query_idx]

        data = list(cypher_results.data())

        if len(data) == 1:
            return {"messages": [AIMessage(content=data)]}

        # 数量大于1的导出到文件中
        # 选取一个文件名
        filename = self.llm.invoke(
            "请根据问题{}和cypher{}生成一个文件名。你只需要输出文件名即可，不需要输出后缀和其他内容。".format(query,
                                                                                                           cypher_query)).content
        # 给文件名加时间戳
        filename = f"{filename}_{int(time.time())}"
        pandas_data = []
        for item in data:
            for k, v in item.items():
                if not isinstance(v, str) and not isinstance(v, int):
                    pandas_data.append(dict(v))
                else:
                    print(v, type(v))

        df = pd.DataFrame(pandas_data)
        df.to_csv(f"output/{filename}.csv", index=False)

        return {
            "messages": [AIMessage(content="文件内容过多，已保存到output目录下，文件名为{}.csv")],
        }

    def draw_graph(self):
        return Image(self.graph.get_graph().draw_png())


if __name__ == '__main__':
    neo4j_agent = Neo4jAgent(checkpointer=memory)
